import asyncio
import threading
import time
from typing import Optional
import signal
import sys
import uvicorn
from datetime import datetime

from config import config
from utils import LogManager, metrics_collector
from cloud_client import cloud_client
from inference_engine import inference_engine
from sample_manager import sample_manager
from edge_server import edge_server

class EdgeMain:
    """边侧主服务"""
    
    def __init__(self):
        # 初始化日志
        LogManager.initialize_logging(config.log.level, config.log.format)
        self.logger = LogManager.get_logger("EdgeMain")
        
        self.running = False
        self.inference_thread = None
        self.server_thread = None
        self.registration_thread = None
        self.health_monitor_thread = None
        
        # 服务状态追踪
        self.node_registered = False
        self.inference_healthy = True
        self.last_registration_attempt = 0
        
        # 注册信号处理
        signal.signal(signal.SIGINT, self._signal_handler)
        signal.signal(signal.SIGTERM, self._signal_handler)
    
    def start(self):
        """启动边侧服务"""
        try:
            self.logger.info("边侧服务启动开始...")
            self.running = True
            
            # 步骤0: 异步注册节点到云端（不阻塞启动）
            self._start_node_registration()
            
            # 检查并加载模型
            self._initialize_model()
            
            # 启动API服务器
            self.logger.info("启动边侧API服务器...")
            self._start_api_server()
            
            # 启动推理定时任务
            self.logger.info("启动推理定时任务...")
            self._start_inference_scheduler()
            
            # 启动服务健康监控
            self.logger.info("启动服务健康监控...")
            self._start_health_monitor()
            
            self.logger.info("边侧服务启动完成！")
            
            # 主循环 - 保持服务运行
            self._main_loop()
            
        except Exception as e:
            self.logger.error(f"边侧服务启动失败: {str(e)}")
            self.stop()
    
    def _start_node_registration(self):
        """启动节点注册线程"""
        def registration_worker():
            """节点注册工作线程"""
            self.logger.info("启动节点注册工作线程")
            retry_interval = 30  # 初始重试间隔30秒
            max_retry_interval = 300  # 最大重试间隔5分钟
            consecutive_failures = 0  # 连续失败次数
            max_consecutive_failures = 10  # 最大连续失败次数
            
            while self.running:
                try:
                    current_time = time.time()
                    
                    # 如果未注册或距离上次尝试超过重试间隔
                    if not self.node_registered and (current_time - self.last_registration_attempt) >= retry_interval:
                        self.logger.info(f"尝试向云端注册节点... (连续失败次数: {consecutive_failures})")
                        self.last_registration_attempt = current_time
                        
                        if cloud_client.register_node():
                            self.node_registered = True
                            retry_interval = 30  # 重置重试间隔
                            consecutive_failures = 0  # 重置失败计数
                            self.logger.info("节点注册成功")
                        else:
                            consecutive_failures += 1
                            # 指数退避，但不超过最大间隔
                            retry_interval = min(retry_interval * 1.5, max_retry_interval)
                            self.logger.warning(f"节点注册失败，{retry_interval}秒后重试 (连续失败: {consecutive_failures}/{max_consecutive_failures})")
                            
                            # 连续失败过多时，降级为更长的重试间隔
                            if consecutive_failures >= max_consecutive_failures:
                                retry_interval = max_retry_interval
                                self.logger.error(f"连续注册失败过多，降级为最大重试间隔: {max_retry_interval}秒")
                    
                    # 如果已注册，定期发送心跳保持连接
                    elif self.node_registered:
                        time.sleep(60)  # 每分钟检查一次连接状态
                        if not self._check_cloud_connection():
                            self.logger.warning("与云端连接丢失，重新进入注册模式")
                            self.node_registered = False
                            retry_interval = 30
                            consecutive_failures = 0  # 重置失败计数
                    
                    time.sleep(5)  # 基础检查间隔
                    
                except Exception as e:
                    self.logger.error(f"节点注册线程异常: {str(e)}")
                    consecutive_failures += 1
                    time.sleep(10)
            
            self.logger.info("节点注册线程停止")
        
        self.registration_thread = threading.Thread(target=registration_worker, daemon=True)
        self.registration_thread.start()
    
    def _check_cloud_connection(self) -> bool:
        """检查与云端的连接状态"""
        try:
            return cloud_client.check_cloud_connection()
        except Exception as e:
            self.logger.debug(f"云端连接检查失败: {str(e)}")
            return False
    
    def _initialize_model(self):
        """初始化模型"""
        self.logger.info("初始化推理模型...")
        
        if not config.inference.auto_load_default:
            self.logger.info("配置为不自动加载模型，推理功能暂时禁用")
            return
        
        # 根据服务类型从云端获取模型
        service_type = config.edge_node.service_type
        self.logger.info(f"根据服务类型 '{service_type}' 获取适用的模型...")
        
        try:
            # 检查云端连接
            if not cloud_client.check_cloud_connection():
                self.logger.warning("无法连接到云端服务器，跳过模型获取")
                self.logger.info("请检查:")
                self.logger.info("   1. 云端服务器是否正在运行")
                self.logger.info("   2. 网络连接是否正常")
                self.logger.info("   3. config.yaml中的云端服务器地址是否正确")
                self.logger.info("推理功能将被禁用，可以稍后通过模型更新接口下发新模型")
                return
            
            # 根据服务类型从云端获取模型
            success, result = cloud_client.get_model_by_service_type(service_type)
            
            if success:
                self.logger.info(f"成功获取适用于 '{service_type}' 的模型: {result}")
                
                # 尝试加载下载的模型
                if inference_engine.load_model(result):
                    self.logger.info("云端模型加载成功，推理功能已启用")
                else:
                    self.logger.error("云端模型加载失败")
            else:
                self.logger.warning(f"获取适用于 '{service_type}' 的模型失败: {result}")
                
                # 提供详细的错误信息和解决方案
                if "没有可用模型" in result:
                    self.logger.info("当前云端没有可用模型，可能的原因:")
                    self.logger.info("   • 云端尚未进行任何训练任务")
                    self.logger.info("   • 模型文件已被删除或移动")
                    self.logger.info("   • 数据库中的模型记录不完整")
                    self.logger.info("")
                    self.logger.info("建议解决方案:")
                    self.logger.info("   1. 在云端创建新的训练任务")
                    self.logger.info("   2. 检查云端models目录中是否有模型文件")
                    self.logger.info("   3. 使用 'curl http://127.0.0.1:8000/api/v1/models' 检查云端模型")
                    self.logger.info(f"   4. 确保云端有适用于 '{service_type}' 的模型")
                    
                self.logger.info("推理功能将被禁用，可以稍后通过模型更新接口下发新模型")
                
        except Exception as e:
            self.logger.error(f"获取模型时发生异常: {str(e)}")
            self.logger.info("推理功能将被禁用，可以稍后通过模型更新接口下发新模型")
    
    def _start_health_monitor(self):
        """启动服务健康监控"""
        def health_monitor():
            """健康监控线程"""
            self.logger.info("服务健康监控启动")
            consecutive_failures = 0
            max_failures = 3
            
            while self.running:
                try:
                    # 检查推理引擎健康状态
                    if inference_engine.has_model():
                        # 尝试进行一次简单的推理测试（可选）
                        self.inference_healthy = True
                        consecutive_failures = 0
                    else:
                        self.logger.debug("推理引擎无模型，尝试重新加载默认模型")
                        if config.inference.auto_load_default:
                            if inference_engine.load_default_model():
                                self.logger.info("健康监控：重新加载默认模型成功")
                                self.inference_healthy = True
                                consecutive_failures = 0
                            else:
                                consecutive_failures += 1
                    
                    # 检查关键目录是否存在
                    self._ensure_critical_directories()
                    
                    # 如果连续失败超过阈值，记录严重错误
                    if consecutive_failures >= max_failures:
                        self.logger.error(f"推理服务连续失败 {consecutive_failures} 次，请检查系统状态")
                        self.inference_healthy = False
                    
                    time.sleep(30)  # 每30秒检查一次
                    
                except Exception as e:
                    self.logger.error(f"健康监控异常: {str(e)}")
                    consecutive_failures += 1
                    time.sleep(10)
            
            self.logger.info("服务健康监控停止")
        
        self.health_monitor_thread = threading.Thread(target=health_monitor, daemon=True)
        self.health_monitor_thread.start()
    
    def _ensure_critical_directories(self):
        """确保关键目录存在"""
        from pathlib import Path
        critical_dirs = [
            Path(config.dataset.temp_dir),
            Path(config.dataset.datasets_dir),
            Path(config.dataset.datasets_dir) / 'images',
            Path(config.dataset.datasets_dir) / 'labels',
            Path(config.dataset.datasets_dir) / config.dataset.confidence_dir,
            Path("models")
        ]
        
        for dir_path in critical_dirs:
            if not dir_path.exists():
                dir_path.mkdir(parents=True, exist_ok=True)
                self.logger.debug(f"重新创建目录: {dir_path}")

    def stop(self):
        """停止边侧服务"""
        self.logger.info("边侧服务停止中...")
        self.running = False
        
        # 停止各个线程
        threads = [self.inference_thread, self.registration_thread, self.health_monitor_thread]
        for thread in threads:
            if thread and thread.is_alive():
                thread.join(timeout=5)
        
        self.logger.info("边侧服务已停止")
    
    def _start_api_server(self):
        """启动API服务器"""
        def run_server():
            try:
                uvicorn.run(
                    edge_server.app,
                    host=config.api.host,
                    port=config.api.port,
                    log_level="warning"  # 降低uvicorn日志级别
                )
            except Exception as e:
                self.logger.error(f"API服务器启动失败: {str(e)}")
        
        self.server_thread = threading.Thread(target=run_server, daemon=True)
        self.server_thread.start()
        
        # 等待服务器启动
        time.sleep(2)
        self.logger.info(f"API服务器已在 {config.api.host}:{config.api.port} 启动")
    
    def _start_inference_scheduler(self):
        """启动推理调度器"""
        def inference_scheduler():
            """推理调度器线程"""
            self.logger.info("推理调度器启动")
            last_status_update = 0
            last_model_check = 0
            
            while self.running:
                try:
                    current_time = time.time()
                    
                    # 步骤1: 自动推理新数据（只有加载模型后才进行推理）
                    if inference_engine.has_model():
                        processed_count = inference_engine.process_temp_images()
                        if processed_count > 0:
                            self.logger.info(f"推理处理完成: {processed_count} 张图片")
                    else:
                        # 如果没有模型，检查是否有新的模型可以加载
                        if current_time - last_model_check >= 300:  # 每5分钟检查一次
                            # 首先尝试加载本地默认模型
                            if inference_engine.load_default_model():
                                self.logger.info("检测到本地新模型，已自动加载")
                            else:
                                # 尝试从云端获取模型
                                self._try_download_model_from_cloud()
                            last_model_check = current_time
                    
                    # 定期清理旧的上传文件和任务记录
                    sample_manager.cleanup_old_uploads()
                    from edge_server import edge_server
                    edge_server.cleanup_old_tasks()
                    
                    # 定期更新节点状态到云端（每分钟一次）
                    if current_time - last_status_update >= config.monitoring.update_interval:
                        self._update_node_status()
                        last_status_update = current_time
                    
                    # 等待下一次检查
                    for _ in range(config.inference.check_interval):
                        if not self.running:
                            break
                        time.sleep(1)
                        
                except Exception as e:
                    self.logger.error(f"推理调度器异常: {str(e)}")
                    time.sleep(5)  # 出错后等待5秒再继续
            
            self.logger.info("推理调度器停止")
        
        self.inference_thread = threading.Thread(target=inference_scheduler, daemon=True)
        self.inference_thread.start()
    
    def _update_node_status(self):
        """更新节点状态到云端"""
        if not self.node_registered:
            self.logger.debug("节点未注册，跳过状态更新")
            return
            
        try:
            # 获取系统指标
            system_metrics = metrics_collector.get_system_metrics()
            
            # 获取样本统计
            sample_stats = sample_manager.get_sample_statistics()
            
            # 获取当前实际部署的模型列表
            deployed_models = []
            if inference_engine.has_model():
                # 优先使用推理引擎中的模型ID
                if hasattr(inference_engine, 'model_metadata') and inference_engine.model_metadata:
                    model_id = inference_engine.model_metadata.get("model_id")
                    if model_id:
                        deployed_models.append(model_id)
                # 如果没有model_id，则使用配置中的模型名作为fallback
                elif config.inference.model_name:
                    deployed_models.append(config.inference.model_name)
            
            # 准备状态数据
            status_data = {
                "status": "online" if self.inference_healthy else "degraded",
                "cpu_load": system_metrics.get("cpu_load", 0.0),
                "memory_load": system_metrics.get("memory_load", 0.0),
                "network_load": system_metrics.get("network_load", 0.0),
                "collected_sample_count": sample_stats.get("collected_samples", 0),
                "uncollected_sample_count": sample_stats.get("uncollected_samples", 0),
                "running_services": [config.edge_node.service_type],
                "deployed_models": deployed_models,
                "service_health": {
                    "inference_engine": self.inference_healthy,
                    "node_registered": self.node_registered
                }
            }
            
            # 更新到云端
            if cloud_client.update_node_status(status_data):
                self.logger.debug("节点状态更新成功")
            else:
                self.logger.debug("节点状态更新失败，可能需要重新注册")
                self.node_registered = False
                
        except Exception as e:
            self.logger.debug(f"更新节点状态异常: {str(e)}")
            # 连接异常可能意味着需要重新注册
            self.node_registered = False
    
    def _try_download_model_from_cloud(self):
        """尝试从云端下载模型"""
        try:
            if not self.node_registered:
                self.logger.debug("节点未注册，跳过云端模型检查")
                return
                
            self.logger.debug("🔍 检查云端是否有可用模型...")
            
            # 从云端获取并下载默认模型
            success, result = cloud_client.get_default_model_from_cloud()
            
            if success:
                self.logger.info(f"✅ 从云端下载模型成功: {result}")
                
                # 尝试加载下载的模型
                if inference_engine.load_model(result):
                    self.logger.info("云端下载的模型加载成功，推理功能已启用")
                    self.inference_healthy = True
                else:
                    self.logger.error("云端下载的模型加载失败")
            else:
                if "没有可用模型" in result:
                    self.logger.debug("云端暂时没有可用模型，等待下次检查")
                else:
                    self.logger.debug(f"暂时无法从云端获取模型: {result}")
                
        except Exception as e:
            self.logger.debug(f"检查云端模型时发生异常: {str(e)}")
    
    def _main_loop(self):
        """主循环"""
        try:
            while self.running:
                time.sleep(1)
        except KeyboardInterrupt:
            self.logger.info("收到停止信号")
        finally:
            self.stop()
    
    def _signal_handler(self, signum, frame):
        """信号处理器"""
        self.logger.info(f"收到信号 {signum}，准备停止服务...")
        self.running = False
    
    def get_service_info(self) -> dict:
        """获取服务信息"""
        return {
            "service_name": "边侧节点服务",
            "node_id": config.edge_node.id,
            "node_name": config.edge_node.name,
            "location": config.edge_node.location,
            "api_port": config.api.port,
            "running": self.running,
            "model_loaded": inference_engine.has_model(),
            "cloud_connected": cloud_client.check_cloud_connection(),
            "start_time": datetime.now().isoformat()
        }

def main():
    """主函数"""
    print("=" * 60)
    print("边侧节点服务启动")
    print("=" * 60)
    
    # 创建并启动边侧服务
    edge_main = EdgeMain()
    try:
        edge_main.start()
    except KeyboardInterrupt:
        print("\n收到中断信号，正在关闭服务...")
    except Exception as e:
        print(f"服务异常退出: {e}")
    finally:
        edge_main.stop()

if __name__ == "__main__":
    main() 